# -*- coding: utf-8 -*-
"""
Created on Fri Jul 12 18:00:17 2013

@author: Maxim
"""

#import sys
from sys import argv
from sys import exit as _exit_
from PyQt4 import QtGui
from main_window import Ui_MainWindow
from part_designer import CadDesigner
from numpy import array, zeros, vstack, insert, flipud
from matplotlib.colors import normalize
from os import getcwd
from os.path import abspath, dirname

from airfoil_widget import AirfoilWidget
from nose_widget import NoseWidget, NoseShape

class DesignerMain():
    """
    Main class of configuration designer.
    """
    def __init__(self):
        self.pd = CadDesigner()
        app = QtGui.QApplication(argv)
        self.mainWindow = QtGui.QMainWindow()
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self.mainWindow)
        self.initialize_buttons()
        self.mainWindow.show()
        self.previousDirectory = getcwd()
        self.airfoil_dialog = AirfoilWidget()
        self.nose_dialog = NoseWidget()
        _exit_(app.exec_())
        
    def initialize_buttons(self):
        self.isTail = False
        self.isRoot = True
        # general
        self.ui.pushButton.clicked.connect(self.select_workdir)

        # wing tab
        self.ui.pushButton_2.clicked.connect(self.add_new_part)
        self.ui.pushButton_3.clicked.connect(self.open_airfoil_dialog)
        self.ui.pushButton_5.clicked.connect(self.open_airfoil_dialog)
        self.ui.pushButton_6.clicked.connect(self.create_wing_section)
        self.ui.pushButton_48.clicked.connect(self.pd.close_wing_tip)
        self.ui.pushButton_7.clicked.connect(self.create_x_tail)
        self.ui.pushButton_8.clicked.connect(self.rotate_wing)
        self.ui.pushButton_9.clicked.connect(self.pd.symmetry_wing)
        self.ui.pushButton_10.clicked.connect(self.save_wing_part)
        self.ui.pushButton_11.clicked.connect(self.save_part_as)
        self.ui.pushButton_12.clicked.connect(lambda: self.close_part('wing'))

        # body tab
        self.ui.pushButton_4.clicked.connect(self.pd.add_new_part)
        self.ui.pushButton_15.clicked.connect(self.save_body_part)
        self.ui.pushButton_16.clicked.connect(self.save_part_as)
        self.ui.pushButton_17.clicked.connect(lambda: self.close_part('body'))
        self.ui.pushButton_13.clicked.connect(self.create_body)
        self.ui.pushButton_44.clicked.connect(self.create_canopy)
        self.ui.pushButton_45.clicked.connect(self.reload_body)
        self.ui.pushButton_46.clicked.connect(lambda: self.open_nose_dialog(False))
        self.ui.pushButton_47.clicked.connect(lambda: self.open_nose_dialog(True))
        # assembly tab
        self.ui.pushButton_14.clicked.connect(self.pd.add_new_product)
        self.ui.pushButton_19.clicked.connect(lambda: self.add_part_to_product(1))
        self.ui.pushButton_21.clicked.connect(lambda: self.add_part_to_product(2))
        self.ui.pushButton_22.clicked.connect(lambda: self.add_part_to_product(3))
        self.ui.pushButton_23.clicked.connect(lambda: self.add_part_to_product(4))
        self.ui.pushButton_24.clicked.connect(lambda: self.add_part_to_product(5))
        self.ui.pushButton_25.clicked.connect(lambda: self.add_part_to_product(6))
        self.ui.pushButton_26.clicked.connect(lambda: self.add_part_to_product(7))
        self.ui.pushButton_27.clicked.connect(lambda: self.add_part_to_product(8))
        self.ui.pushButton_28.clicked.connect(lambda: self.add_part_to_product(9))
        self.ui.pushButton_29.clicked.connect(lambda: self.add_part_to_product(10))
        self.ui.pushButton_30.clicked.connect(lambda: self.add_part_to_product(11))
        self.ui.pushButton_20.clicked.connect(lambda: self.move_part(1))
        self.ui.pushButton_31.clicked.connect(lambda: self.move_part(2))
        self.ui.pushButton_32.clicked.connect(lambda: self.move_part(3))
        self.ui.pushButton_33.clicked.connect(lambda: self.move_part(4))
        self.ui.pushButton_34.clicked.connect(lambda: self.move_part(5))
        self.ui.pushButton_35.clicked.connect(lambda: self.move_part(6))
        self.ui.pushButton_36.clicked.connect(lambda: self.move_part(7))
        self.ui.pushButton_37.clicked.connect(lambda: self.move_part(8))
        self.ui.pushButton_38.clicked.connect(lambda: self.move_part(9))
        self.ui.pushButton_39.clicked.connect(lambda: self.move_part(10))
        self.ui.pushButton_40.clicked.connect(lambda: self.move_part(11))
        self.ui.pushButton_41.clicked.connect(self.save_product)
        self.ui.pushButton_42.clicked.connect(self.save_product_as)
        self.ui.pushButton_43.clicked.connect(self.close_product)
    
    def _init_wing_buttons(self):
        self.ui.pushButton_3.setEnabled(False)
        self.ui.pushButton_5.setEnabled(False)
        self.ui.pushButton_6.setEnabled(False)
        self.ui.pushButton_7.setEnabled(False)
        self.ui.pushButton_8.setEnabled(False)
        self.ui.pushButton_9.setEnabled(False)
        self.ui.pushButton_10.setEnabled(False)
        self.ui.pushButton_11.setEnabled(False)
        self.ui.pushButton_12.setEnabled(False)
    
    def _init_body_buttons(self):
        self.ui.pushButton_4.setEnabled(True)
        self.ui.pushButton_13.setEnabled(False)
        self.ui.pushButton_15.setEnabled(False)
        self.ui.pushButton_16.setEnabled(False)
        self.ui.pushButton_17.setEnabled(False)
    
    def _init_assembly_buttons(self):
        self.ui.pushButton_14.setEnabled(True)
        self.ui.pushButton_43.setEnabled(False)
        self.ui.pushButton_19.setEnabled(False)
        self.ui.pushButton_21.setEnabled(False)
        self.ui.pushButton_22.setEnabled(False)
        self.ui.pushButton_23.setEnabled(False)
        self.ui.pushButton_24.setEnabled(False)
        self.ui.pushButton_25.setEnabled(False)
        self.ui.pushButton_26.setEnabled(False)
        self.ui.pushButton_27.setEnabled(False)
        self.ui.pushButton_28.setEnabled(False)
        self.ui.pushButton_29.setEnabled(False)
        self.ui.pushButton_30.setEnabled(False)
        self.ui.pushButton_20.setEnabled(False)
        self.ui.pushButton_31.setEnabled(False)
        self.ui.pushButton_32.setEnabled(False)
        self.ui.pushButton_33.setEnabled(False)
        self.ui.pushButton_34.setEnabled(False)
        self.ui.pushButton_35.setEnabled(False)
        self.ui.pushButton_36.setEnabled(False)
        self.ui.pushButton_37.setEnabled(False)
        self.ui.pushButton_38.setEnabled(False)
        self.ui.pushButton_39.setEnabled(False)
        self.ui.pushButton_40.setEnabled(False)
        self.ui.pushButton_41.setEnabled(False)
        self.ui.pushButton_42.setEnabled(False)
    
    def add_new_part(self):
        self.pd.add_new_part()
        self.isRoot = True

    def close_part(self,part):
        if part=='wing':
            self._init_wing_buttons()
            self.pd.close_part()
        elif part=='body':
            self._init_body_buttons()
            self.pd.close_part()
    
    def close_product(self):
        self._init_assembly_buttons()
        self.pd.close_product()

    def create_canopy(self):
        xref = float(self.ui.lineEdit_11.text()) *1.0e3
        yref = 0.0
        zref = float(self.ui.lineEdit_13.text()) *1.0e3
        refPt = array([xref,yref,zref])
        upCurve     = self._get_coord_from_raw(self.ui.plainTextEdit_2.toPlainText()) * 1.0e3
        upCurve = insert(upCurve,1,zeros(upCurve.shape[0]),axis=1)
        sideCurve   = self._get_coord_from_raw(self.ui.plainTextEdit_3.toPlainText()) * 1.0e3
        sideCurve = insert(sideCurve,2,zeros(sideCurve.shape[0]),axis=1)
        csLocation  = self._get_coord_from_raw(self.ui.plainTextEdit_4.toPlainText()) / 100.0
        csConic     = self._get_coord_from_raw(self.ui.plainTextEdit_5.toPlainText())
        sideCurve  = vstack([upCurve[0],sideCurve,upCurve[-1]])
        upCurve   += refPt
        sideCurve += refPt
        norm1      = normalize(upCurve[0,0],upCurve[-1,0])
        csLocation = array([norm1.inverse(val) for val in csLocation])
        self.pd.create_canopy(upCurve,sideCurve,csLocation,csConic)
        
    def reload_body(self):
        self.close_part('body')
        self.add_new_part()
        self.create_body()
        self.create_canopy()
        self.ui.pushButton_15.setEnabled(True)
        self.ui.pushButton_16.setEnabled(True)
        self.ui.pushButton_17.setEnabled(True)

    def create_wing_section(self):
        """
        Creates root airfoil section or tip airfoil and wing surface. Tip 
        section can be controlled by chord or taper ratio.
        """
        if self.isRoot:
            self.isRoot = False
            xref = float(self.ui.lineEdit_3.text()) *1.0e3
            yref = float(self.ui.lineEdit_4.text()) *1.0e3
            zref = float(self.ui.lineEdit_5.text()) *1.0e3
            try:
                self.pd.refPoint = array([xref,yref,zref])
            except ValueError:
                self.pd.refPoint = zeros(3)
            self.chord     = float(self.ui.lineEdit_7.text()) *1.0e3
            self.incidence = float(self.ui.lineEdit_8.text())
            rotAxis        = float(self.ui.lineEdit_9.text()) / 100.0
            af = self._adjust_trailing_edge(self.airfoil_dialog.af, 
                                            self.airfoil_dialog.trailingEdgeThickness, self.chord)
            #print self.airfoil_dialog.isSpline
            self.pd.create_airfoil(af,self.chord,self.incidence,rotAxis,self.airfoil_dialog.isSpline)
        else:
            if self.ui.radioButton.isChecked():
                self.chord = float(self.ui.lineEdit_16.text()) * 1.0e3
            else:
                self.chord = self.chord * float(self.ui.lineEdit_15.text())
            self.incidence += float(self.ui.lineEdit_17.text())
            rotAxis  = float(self.ui.lineEdit_9.text()) / 100.0
            span     = float(self.ui.lineEdit_19.text()) * 1.0e3
            if self.ui.radioButton_3.isChecked():
                sweep  = float(self.ui.lineEdit_20.text())
                c4sweep = False
            elif self.ui.radioButton_4.isChecked():
                sweep = float(self.ui.lineEdit_21.text())
                c4sweep = True
            dihedral = float(self.ui.lineEdit_22.text())
            af = self._adjust_trailing_edge(self.airfoil_dialog.af, 
                                            self.airfoil_dialog.trailingEdgeThickness, self.chord)
            self.pd.create_tip_section(af,self.chord,
                                       self.incidence,rotAxis,span,sweep,dihedral,c4sweep,self.airfoil_dialog.isSpline)
    
    def _adjust_trailing_edge(self,airfoil,trailingEdgeThickness,chord):
        if not trailingEdgeThickness==None:
            if trailingEdgeThickness[0]==0:
                zTE = trailingEdgeThickness[1] / chord
            elif trailingEdgeThickness[0]==1:
                zTE = trailingEdgeThickness[1]
            airfoil.set_trailing_edge_gap(zTE)
        return airfoil

    def create_x_tail(self):
        """
        Creates X-tail configuration. Input is value of first segment angle 
        wrt. X-axis. The rest of segments will be located with 90 degree 
        increment to generate X-tail.
        """
        self.isTail = True
        angle = float(self.ui.lineEdit_23.text())
        self.pd.join_wing(True)
        self.pd.create_xtail(angle)
        self.pd.hide_curves()
        self.pd.hide_surfaces()
    
    def rotate_wing(self):
        """
        Rotates wing along Z-axis with predefined axis root point in XY and angle. 
        This function should be used to create foldable wing configurations at 
        different angles.
        """
        angle = float(self.ui.lineEdit_24.text())
        x = float(self.ui.lineEdit_25.text()) * 1.0e3
        y = float(self.ui.lineEdit_26.text()) * 1.0e3
        self.pd.join_wing(True)
        self.pd.rotate_wing(angle,[x,y])
        self.pd.hide_curves()

    def save_wing_part(self):
        if self.isTail:
            partName = '\\wing.CATPart'
        else:
            partName = '\\tail.CATPart'
        self.pd.save_part(self.ui.lineEdit.text() + partName)
        
    def save_body_part(self):
        path = self.ui.lineEdit.text() + '\\body.CATPart'
        self.pd.save_part(path)
    
    def save_product(self):
        path = self.ui.lineEdit.text() + '\\assembly.CATProduct'
        self.pd.save_product(path)

    def save_part_as(self):
        path = QtGui.QFileDialog.getSaveFileName(self.mainWindow,"Save part",self.previousDirectory,'CATIA Part File (*.CATPart)')
        self.previousDirectory = dirname(abspath(path))
        print self.previousDirectory
        path = abspath(path)
        self.pd.save_part(path)

    def save_product_as(self):
        path = QtGui.QFileDialog.getSaveFileName(self.mainWindow,"Save product",self.previousDirectory,'CATIA Product File (*.CATProduct)')
        path = abspath(path)
        self.previousDirectory = None
        self.pd.save_product(path)

    def select_workdir(self):
        newPath = QtGui.QFileDialog.getExistingDirectory()
        if not newPath=='':
            self.ui.lineEdit.setText(newPath)
    
    def create_body(self):
        # --- read and process data
        nose = self.nose_dialog.nose
        tail = self.nose_dialog.tail
        L1        = float(self.ui.lineEdit_27.text()) *1.0e3 #nose length, mm
        D1        = float(self.ui.lineEdit_29.text()) *1.0e3 #nose diameter, mm
        offset    = float(self.ui.lineEdit_32.text()) *1.0e3 # tail offset, mm
        if self.ui.checkBox.isChecked() or tail==None:
            tail = NoseShape()
            tail.type = 0
            L2 = 0.0
            ptEnd = 0.0
            contEnd = 0
        else:
            L2        = float(self.ui.lineEdit_31.text()) *1.0e3 #tail length, mm
            ptEnd     = 1.0-float(self.ui.lineEdit_67.text())/100.
            contEnd = int(self.ui.comboBox_2.currentIndex())
        if self.ui.checkBox_2.isChecked():
            D2 = D1
        else:
            D2        = float(self.ui.lineEdit_30.text()) *1.0e3 #tail diameter, mm
        ptStart   = float(self.ui.lineEdit_33.text())/100.
        contStart = int(self.ui.comboBox.currentIndex())
        try:
            _coord = self._get_coord_from_raw(self.ui.plainTextEdit.toPlainText()) *1.0e3
            sideProfilePts = zeros([_coord.shape[0]+2,3])
            sideProfilePts[1:-1,0] = _coord[:,0]
            sideProfilePts[1:-1,2] = _coord[:,1]
        except:
            sideProfilePts = None
        #--- actions ---
        noseProfile = self.create_nose_profile(nose,L1,D1)
        nosePt = self.pd.create_pt_on_curve_percentage(noseProfile,ptStart,show=False)

        tailProfile = self.create_nose_profile(tail, L2, D2, offset+L1, True)
        tailPt = self.pd.create_pt_on_curve_percentage(tailProfile,ptEnd,show=False)
        tangency = [None,None]
        curvature = [None,None]
        if contStart==1:
            tangency[0] = noseProfile
        elif contStart==2:
            curvature[0] = noseProfile
        if contEnd==1 and self.ui.checkBox.isChecked()==False:
            tangency[1] = tailProfile
        elif contEnd==2 and self.ui.checkBox.isChecked()==False:
            curvature[1] = tailProfile
        if 0<ptStart<1.0:
            split1 = self.pd.split_curve(noseProfile,nosePt,show=True)
        else:
            split1 = noseProfile
        if 0<ptEnd<1.0:
            split2 = self.pd.split_curve(tailProfile,tailPt,1,show=True)
        else:
            split2 = tailProfile
        refSideProfile = self.pd.create_spline(sideProfilePts,tangency,curvature,pointRef=[nosePt,tailPt])[2]
        profile = self.pd.create_join([split1,refSideProfile,split2],show=True)
        self.pd.create_revolute_x(profile,show=True)
        
    def create_nose_profile(self,profile,length,diameter,offset=0.0,reverse=False,show=False):
        L = length
        D = diameter
        if profile.type==0:
            pts = profile(L,D)
            if reverse:
                pts = self.reverse_points(pts,length,offset)
            if profile.blunt=='spherical':
                refProfile = self.pd.nose_conic_blunt(pts,profile.Rnose,reverse)
            elif profile.blunt=='flat':
                refProfile = self.pd.nose_conic_flat(pts,reverse)
            else:
                refProfile = self.pd.nose_conic(pts,reverse)
        elif profile.type==1:
            pts, rho = profile(L,D)
            if reverse:
                pts = self.reverse_points(pts,length,offset)
            if profile.blunt=='spherical':
                refProfile = self.pd.nose_ogival_blunt(pts,rho,profile.Rnose,reverse)
            elif profile.blunt=='flat':
                refProfile = self.pd.nose_ogival_flat(pts,rho,reverse)
            else:
                refProfile = self.pd.nose_ogival(pts,rho)
        elif profile.type==2:
            pts,r = profile(L,D)
            if reverse:
                pts = self.reverse_points(pts,length,offset)
            refProfile = self.pd.nose_ogival(pts,r)
        else:
            pts = profile(L,D)
            if reverse:
                pts = self.reverse_points(pts,length,offset)
            refProfile = self.pd.create_spline(pts)[2]
        return refProfile

    def reverse_points(self,pts,length,offset):
        pts[:,0] = -pts[:,0]+length+offset
        return flipud(pts)
    
    def create_body_old(self):
        radiusNose  = float(self.ui.lineEdit_29.text()) * 1.0e3
        angleNose   = float(self.ui.lineEdit_30.text())
        if self.ui.radioButton_5.isChecked():
            contNose = 'C0'
        elif self.ui.radioButton_6.isChecked():
            contNose = 'C1'
        elif self.ui.radioButton_7.isChecked():
            contNose = 'C2'
        xoffset    = float(self.ui.lineEdit_33.text()) * 1.0e3
        radiusTail = float(self.ui.lineEdit_32.text()) * 1.0e3
        angleTail  = float(self.ui.lineEdit_31.text())
        if self.ui.radioButton_9.isChecked():
            contTail = 'C0'
        elif self.ui.radioButton_8.isChecked():
            contTail = 'C1'
        elif self.ui.radioButton_10.isChecked():
            contTail = 'C2'
        try:
            _coord = self._get_coord_from_raw(self.ui.plainTextEdit.toPlainText()) * 1.0e3
            coord = zeros([_coord.shape[0],3])
            coord[:,0] = _coord[:,0]
            coord[:,2] = _coord[:,1]
        except:
            coord = None
        self.pd.create_body(radiusNose,angleNose,contNose,radiusTail,angleTail,
                            contTail,xoffset,coord)
    def _get_coord_from_raw(self,string):
        string = str(string)
        for i,line in enumerate(string.splitlines()):
            seg = line.split()
            if i==0:
                data = zeros([30,len(seg)])
            for j,val in enumerate(seg):
                data[i,j] = float(val)
        return data[:i+1]

    def open_airfoil_dialog(self):
        self.airfoil_dialog.widget.show()
    def open_nose_dialog(self,isTail):
        self.nose_dialog.widget.show()
        self.nose_dialog.isTail = isTail
        
    def add_part_to_product(self,n):
        path = QtGui.QFileDialog.getOpenFileName()
        if path!='':
            self.pd.add_part_to_assembly(path,n)

    def move_part(self,n):
        newLocation = zeros(3)
        try:
            if n==1:
                newLocation[0] = float(self.ui.lineEdit_34.text()) * 1e3
                newLocation[1] = float(self.ui.lineEdit_35.text()) * 1e3
                newLocation[2] = float(self.ui.lineEdit_36.text()) * 1e3
            elif n==2:
                newLocation[0] = float(self.ui.lineEdit_37.text()) * 1e3
                newLocation[1] = float(self.ui.lineEdit_38.text()) * 1e3
                newLocation[2] = float(self.ui.lineEdit_39.text()) * 1e3
            elif n==3:
                newLocation[0] = float(self.ui.lineEdit_40.text()) * 1e3
                newLocation[1] = float(self.ui.lineEdit_41.text()) * 1e3
                newLocation[2] = float(self.ui.lineEdit_42.text()) * 1e3
            elif n==4:
                newLocation[0] = float(self.ui.lineEdit_43.text()) * 1e3
                newLocation[1] = float(self.ui.lineEdit_44.text()) * 1e3
                newLocation[2] = float(self.ui.lineEdit_45.text()) * 1e3
            elif n==5:
                newLocation[0] = float(self.ui.lineEdit_46.text()) * 1e3
                newLocation[1] = float(self.ui.lineEdit_47.text()) * 1e3
                newLocation[2] = float(self.ui.lineEdit_48.text()) * 1e3
            elif n==6:
                newLocation[0] = float(self.ui.lineEdit_49.text()) * 1e3
                newLocation[1] = float(self.ui.lineEdit_50.text()) * 1e3
                newLocation[2] = float(self.ui.lineEdit_51.text()) * 1e3
            elif n==7:
                newLocation[0] = float(self.ui.lineEdit_52.text()) * 1e3
                newLocation[1] = float(self.ui.lineEdit_53.text()) * 1e3
                newLocation[2] = float(self.ui.lineEdit_54.text()) * 1e3
            elif n==8:
                newLocation[0] = float(self.ui.lineEdit_55.text()) * 1e3
                newLocation[1] = float(self.ui.lineEdit_56.text()) * 1e3
                newLocation[2] = float(self.ui.lineEdit_57.text()) * 1e3
            elif n==9:
                newLocation[0] = float(self.ui.lineEdit_58.text()) * 1e3
                newLocation[1] = float(self.ui.lineEdit_59.text()) * 1e3
                newLocation[2] = float(self.ui.lineEdit_60.text()) * 1e3
            elif n==10:
                newLocation[0] = float(self.ui.lineEdit_61.text()) * 1e3
                newLocation[1] = float(self.ui.lineEdit_62.text()) * 1e3
                newLocation[2] = float(self.ui.lineEdit_63.text()) * 1e3
            elif n==11:
                newLocation[0] = float(self.ui.lineEdit_64.text()) * 1e3
                newLocation[1] = float(self.ui.lineEdit_65.text()) * 1e3
                newLocation[2] = float(self.ui.lineEdit_66.text()) * 1e3
        except ValueError:
            newLocation = zeros(3)
        self.pd.move_part(n,newLocation)


def main():
    mainForm = DesignerMain()

if __name__ == '__main__':
    main()